package com.javasm.mingming.socket;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.javasm.mingming.chat.Agreement;
import com.javasm.mingming.common.utils.TimeUtil;
import io.reactivex.Flowable;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.Arrays;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author javasm
 * @version 0.1
 * @className AISocket
 * @descriptioin:
 * @date 2025/1/21 9:59
 * @since jdk11
 */
@Component
@ServerEndpoint("/ai/{uid}")
public class AISocket {
    Logger logger = LogManager.getLogger(AISocket.class);
    //记录 当前 和 AI 聊天的 Session信息,Key 是uid value是Session对象
    private static Map<String, Session> userSessionMap = new ConcurrentHashMap<>(8);

    @OnOpen
    public void start(@PathParam("uid") String uid, Session session) {
        userSessionMap.put(uid, session);
    }

    @OnClose
    public void end(@PathParam("uid") String uid, Session session) {
        userSessionMap.remove(uid);
    }

    @OnError
    public void error(Throwable e) {
        e.printStackTrace();
    }

    @OnMessage
    public void replyMessage(@PathParam("uid") String uid, Session session, String message) throws IOException, NoApiKeyException, InputRequiredException {
        //消息:协议,时间,uid,消息

        //自己 收到 自己的消息
        String msg = "%s,%s,%s,%s";
        //Agreement.AiAsk+","+time+","+uid+","+message
        session.getBasicRemote().sendText(String.format(msg, Agreement.AiAsk, TimeUtil.getNow12(), uid, message));
        //提问
        //告诉客户端,接下来,是回答的消息 开始
        //AiAnswer   协议,时间,状态(0在发送,1开始,2结束),消息
        session.getBasicRemote().sendText(String.format(msg, Agreement.AiAnswer, TimeUtil.getNow12(), 1, ""));
        //AI 官方代码
        Generation gen = new Generation();
        //AI的人设,可以指定是一个AI助手
        Message systemMsg = Message.builder()
                .role(Role.SYSTEM.getValue())
                .content("You are a helpful assistant.")
                .build();
        //提问的内容
        Message userMsg = Message.builder()
                .role(Role.USER.getValue())
                .content(message)//向AI 提问
                .build();
        GenerationParam param = GenerationParam.builder()
                // 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey(System.getenv("DASHSCOPE_API_KEY"))
                // 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
                .model("qwen-plus")
                .messages(Arrays.asList(systemMsg, userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .topP(0.8).enableSearch(true)//设置 流输出
                .incrementalOutput(true)//流 式 输出
                .build();
        Flowable<GenerationResult> result = gen.streamCall(param);
        //回答,循环回答
        result.blockingForEach(aimessage -> {
            String aiAnswer = aimessage.getOutput().getChoices().get(0).getMessage().getContent();
            session.getBasicRemote().sendText(String.format(msg, Agreement.AiAnswer, TimeUtil.getNow12(), "0", aiAnswer));
        });
        //回答结束
        session.getBasicRemote().sendText(String.format(msg, Agreement.AiAnswer, TimeUtil.getNow12(), "2", ""));
    }
}
